package com.liuhengjia.aiknowledgebase.advisor;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.liuhengjia.aiknowledgebase.common.MessageWrapper;
import com.liuhengjia.aiknowledgebase.entity.BaseConversation;
import com.liuhengjia.aiknowledgebase.service.BaseConversationService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 基于数据库的会话记忆功能
 */
@Service
@RequiredArgsConstructor
public class DbChatMemory implements ChatMemory {
    private final BaseConversationService baseConversationService;

    /**
     * 将会话信息添加到数据库
     *
     * @param conversationId 会话id
     * @param messageList    消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messageList) {
        BaseConversation baseConversation = baseConversationService.getOneOpt(
                        new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId))
                .orElseThrow(
                        () -> new RuntimeException("can not find the conversation , conversation id is %s".formatted(conversationId)));
        List<Message> savedMessages = convert(baseConversation);
        savedMessages.addAll(messageList);
        baseConversationService.update(BaseConversation.builder()
                        .content(MessageWrapper.toConversationStr(savedMessages.stream().map(MessageWrapper::new).toList())).build(),
                new LambdaQueryWrapper<BaseConversation>()
                        .eq(BaseConversation::getConversationId, conversationId));
    }

    /**
     * 从数据库提取会话信息
     *
     * @param conversationId 会话id
     * @param lastN          后几条
     * @return 消息列表
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        return baseConversationService.getOneOpt(new LambdaQueryWrapper<BaseConversation>()
                        .eq(BaseConversation::getConversationId, conversationId))
                .map(v -> convert(v).stream().limit(lastN).toList())
                .orElse(List.of());
    }

    /**
     * 清理会话内容
     *
     * @param conversationId 会话id
     */
    @Override
    public void clear(String conversationId) {
        baseConversationService.remove(
                new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId));
    }

    /**
     * Model 转会话列表
     *
     * @param conversation 会话id
     * @return 会话列表
     */
    private List<Message> convert(BaseConversation conversation) {
        if (!StringUtils.hasText(conversation.getContent())) {
            return new ArrayList<>();
        }
        return MessageWrapper.fromConversationStr(conversation.getContent()).stream().map(MessageWrapper::toMessage).collect(
                Collectors.toList());
    }
}
